OpenAIのEmbeddings APIのベクトルを使って検索を行う

今回はOpenAIのEmbeddings APIを利用して文章検索を行ってみます。このAPIを利用することで文章をベクトルに変換することが可能です。この変換されたベクトル間の距離を計算することで、関連する文章が計算できるようになります。

#OpenAI

#Python

アシ

2023.05.02

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

OpenAIのEmbeddingについて

OpenAIのAPIの一つにEmbeddings APIというものがあります。

これを利用すると文章をベクトルに変換することが可能です。この変換されたベクトルは以下のような用途で利用できます。

検索（クエリとの関連性に基づくランキングの作成）
クラスタリング（文章の類似性によるグループ化）
レコメンデーション（関連する文章を持つ物のレコメンデーション）
異常検出（関連性が低い外れ値の特定）
多様性測定（類似性の分布の分析）
分類（最も類似したラベルでの分類）

ここで言うベクトル形式への変換は以下のようなイメージです。

人間にとっては意味不明ですが、生のテキストデータと比べてコンピュータで扱いやすくなっています。ベクトル化することで様々な計算アルゴリズムが適用できるようになるので、便利です。

また、LLMを使用することで、近い文章は近い位置にデータが配置されるようになります。これによって表現の揺れにも強くなります。検索を例に取ると以下のようなイメージです。

今回はこのEmbeddings APIを利用して検索システムのプロトタイプを作っていきます。

インデックスの作成

まずはじめに、検索の対象となるデータベースを作っていきます。

今回は以下のような文章を１０個用意しました。文章自体はChatGPTに作ってもらいました（コードは付録に書いておきます。）

タイトル: パイナップル
本文: パイナップルは、南アメリカ原産の常緑性の熱帯果樹で、アブラナ科の植物です。果肉は黄色く、鮮やかな甘酸っぱい香りと味わいを持ち、豊富なビタミンCやカロテン、ポリフェノール、カルシウムなどの栄養素が含まれています。また、消化酵素であるブロメリンを含んでおり、食欲増進や消化促進効果があるとされています。切り方や調理法によっては、サラダやスムージー、パイやジュースなど、幅広い料理に使用されます。しかし、果肉とともに硬い芯があるため、適切な切り方をしなければ喉に詰まることがあるため注意が必要です。

パイナップルはアブラナ科ではないようなので、この文章自体は誤りが含まれていますが、検索には影響無いはずなのでこのまま使います。

パイナップル - Wikipedia

この文章をベクトル化しデータベースを作っていきます。今回はシンプルにJSONにすべて保存します。 embeddingフィールドにEmbeddingを利用して出力されたベクトルが格納されてます。

[
    {
        "title": "パイナップル",
        "body": "パイナップルは、(以下略)"
        "embedding": [0.12345, ..., 以下略]
    },
    {
        "title": "Python",
        "body": "Pythonは、(以下略)"
        "embedding": [0.678901, ..., 以下略]
    }
]

コードの全文は以下のとおりです。

import openai
import os

import json

openai.api_key = os.environ['OAI_API_KEY']

# 入力用の文章をロード
with open('docs.json') as f:
    docs = json.load(f)

index = []
for doc in docs:
    # ここでベクトル化を行う
    # openai.embeddings_utils.embeddings_utilsを使うともっとシンプルにかけます
    res = openai.Embedding.create(
        model='text-embedding-ada-002',
        input=doc['body']
    )

    # ベクトルをデータベースに追加
    index.append({
        'title': doc['title'],
        'body': doc['body'],
        'embedding': res['data'][0]['embedding']
    })

with open('index.json', 'w') as f:
    json.dump(index, f)

今回はモデルとしてtext-embedding-ada-002を使用しています。これが2023/5/2時点での最新のモデルのようで、第一世代モデルと異なり、検索、類似性等の用途を問わず利用できるみたいです。

詳細は以下の記事が参考になるかと思います。

text-embedding-ada-002の概要

価格は1000トークンあたり0.0004ドルです。(※2023/5/2時点) 日本語は１文字で１トークン扱いのはずです。

太宰治の人間失格が78000字程度なので0.03ドルですべてベクトル化できるといった感じです。

検索

つづいて検索部分です。検索のロジックとしては以下のとおりです。

クエリ（検索用の文章）をベクトル化
データベース内のすべての文章とのコサイン類似度を計算
コサイン類似度で降順でソート

2番のコサイン類似度のところは今回はシンプルに総当りで計算しています。

コードは以下のとおりです。

import openai
from openai.embeddings_utils import cosine_similarity
import os

import json

openai.api_key = os.environ['OAI_API_KEY']

# データベースの読み込み
with open('index.json') as f:
    INDEX = json.load(f)

# これが検索用の文字列
QUERY = 'なんか甘くてオレンジのやつ'

# 検索用の文字列をベクトル化
query = openai.Embedding.create(
    model='text-embedding-ada-002',
    input=QUERY
)

query = query['data'][0]['embedding']

# 総当りで類似度を計算
results = map(
        lambda i: {
            'title': i['title'],
            'body': i['body'],
            # ここでクエリと各文章のコサイン類似度を計算
            'similarity': cosine_similarity(i['embedding'], query)
            },
        INDEX
)
# コサイン類似度で降順（大きい順）にソート
results = sorted(results, key=lambda i: i['similarity'], reverse=True)

# 以下で結果を表示
print(f"Query: {QUERY}")
print("Rank: Title Similarity")
for i, result in enumerate(results):
    print(f'{i+1}: {result["title"]} {result["similarity"]}')

print("====Best Doc====")
print(f'title: {results[0]["title"]}')
print(f'body: {results[0]["body"]}')

print("====Worst Doc====")
print(f'title: {results[-1]["title"]}')
print(f'body: {results[-1]["body"]}')

OpenAIの公式ドキュメントで推奨されているとおりに、クエリと文章間の距離はコサイン類似度を使用しています。ただ、どの距離関数を使うかはそこまで重要じゃないみたいですね（ユークリッド距離でも同じ結果になるようです）

Which distance function should I use?

実行した結果は以下のようになります。

Query: なんか甘くてオレンジのやつ
Rank: Title Similarity
1: パイナップル 0.7891167093239757
2: 焼き肉 0.7650508389800532
3: 迷彩柄 0.7507892687728419
4: パトカー 0.7388379877365951
5: Python 0.7336562275814604
6: 成人 0.7314125660237523
7: 挑戦状 0.7293262064799704
8: 竜巻 0.7256960383097651
9: 写真撮影 0.7244692550536843
10: 正式名称 0.7200881381916038
====Best Doc====
title: パイナップル
body: パイナップルは、南アメリカ原産の常緑性の熱帯果樹で、アブラナ科の植物です。果肉は黄色く、鮮やかな甘酸っぱい香りと味わいを持ち、豊富なビタミンCやカロテン、ポリフェノール、カルシウムなどの栄養素が含まれています。また、消化酵素であるブロメリンを含んでおり、食欲増進や消化促進効果があるとされています。切り方や調理法によっては、サラダやスムージー、パイやジュースなど、幅広い料理に使用されます。しかし、果肉とともに硬い芯があるため、適切な切り方をしなければ喉に詰まることがあるため注意が必要です。
====Worst Doc====
title: 正式名称
body: 「正式名称」は、物や人物、団体などに対して公式に決まっている呼び名のことです。正確な名称を使うことで、その物や人物、団体などを明確に区別することができます。例えば、企業の正式名称は、商業登記簿に登録された名称とされます。また、政府の機関の名称は、国や地域によって異なりますが、それぞれ決められた公式の名称が存在します。正式名称は、広報や報道機関などで使われることが多く、また、ビジネスや法律などにおいても重要な役割を果たします。正確な名称を使用することで、情報共有を正確に行い、行政手続きなどでもスムーズなやりとりができるようになります。

今回は「なんか甘くてオレンジのやつ」で検索しましたが、パイナップルが１位になっています。オレンジ色(文章中では黄色)というノイズが乗ってますが、影響を受けず関連性が高く出ています。

２番が焼き肉なのも食べ物なため妥当だと思います。

おわりに

今回はOpenAIのEmbeddings APIを利用して検索システムのプロトタイプを作ってみました。今回は１０個程度の文章だったので、シンプルな方法でも十分高速に検索できましたが、文章量が多い場合は専用のDBを使うなど工夫が必要になるでしょう。

また、Embeddings APIは一度に入力できるトークンの数に上限があるため、長い文章をベクトル化するときにはLangChainなどの外部のパッケージを使うといった工夫が必要になります。

付録

ChatGPTによるダミー文章の作成

import openai
import os

import json

openai.api_key = os.environ['OAI_API_KEY']

titles = [
    'パトカー',
    'Python',
    '写真撮影',
    '正式名称',
    'パイナップル',
    '挑戦状',
    '成人',
    '焼き肉',
    '迷彩柄',
    '竜巻',
]

SYSTEM_PROMPT = '''
提供される単語を300字以内で説明してください。
'''

docs = []
for title in titles:
    res = openai.ChatCompletion.create(
        model='gpt-3.5-turbo',
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": title}
        ]
    )

    docs.append({
        'title': title,
        'body': res.choices[0].message.content
    })
    print(f'タイトル: {title}')
    print(res.choices[0].message.content)

with open('docs.json', 'w') as f:
    json.dump(docs, f)

生成された文章

内容は真偽不明です。

title: パトカー
body: パトカーとは、警察が緊急時や巡回監視などのために使用する車両のことを指します。高速移動や急ブレーキなどの過酷な運転条件に耐えうるよう、耐久性や速度性能などが高く設計されています。また、警察官が緊急時の迅速な出動や現場到着を目的に、赤色や青色の回転灯、サイレンなどを装備しています。一般的なパトカーには、4ドアセダンやSUVなどが使われていますが、中にはハイパフォーマンスカーを使用する警察もあります。パトカーは、社会の安全を守るために欠かせない存在となっており、一般道でも見かけることがあります。

title: Python
body: Pythonは、オープンソースのプログラミング言語で、1991年に発表されました。Pythonは、シンプルで読みやすい文法により、学習が容易であり、豊富なライブラリにより、多種多様な分野で利用されています。また、Pythonはフルスタックのウェブアプリケーション開発、データサイエンス、機械学習、人工知能の開発、自然言語処理、画像処理、ブロックチェーンなどの分野で広く使われています。Pythonは対話型モード、スクリプトモード、関数型プログラミング、オブジェクト指向プログラミングなど、多岐にわたるプログラミングスタイルをサポートしています。Pythonは、Windows、Linux、Mac OS Xなどの多くのプラットフォームで動作します。また、PythonはNumPy、Pandas、Matplotlibなどの多数のライブラリを提供しており、これらは大量のデータを扱えるように設計されています。

title: 写真撮影
body: 写真撮影とは、カメラを使って光を捕捉し、それを記録に残すことによって、現実の瞬間を他の人たちと共有できるようにする行為です。写真撮影には様々な種類があり、ポートレートや風景、スポーツなど、さまざまなシチュエーションで撮影されます。写真撮影には、カメラの種類や撮影技術、照明の知識、ポーズの取り方、画像編集など、多くの要素が含まれます。また、撮影場所や被写体の性格や雰囲気など、実際の現場での対応力も重要な要素のひとつです。最近では、スマートフォンによる写真撮影も一般的になっており、誰でも手軽に写真を撮ることができるようになっています。写真撮影は、美術や広告など、様々な分野で利用されており、ビジネスの一部としても重要な役割を果たしています。

title: 正式名称
body: 「正式名称」は、物や人物、団体などに対して公式に決まっている呼び名のことです。正確な名称を使うことで、その物や人物、団体などを明確に区別することができます。例えば、企業の正式名称は、商業登記簿に登録された名称とされます。また、政府の機関の名称は、国や地域によって異なりますが、それぞれ決められた公式の名称が存在します。正式名称は、広報や報道機関などで使われることが多く、また、ビジネスや法律などにおいても重要な役割を果たします。正確な名称を使用することで、情報共有を正確に行い、行政手続きなどでもスムーズなやりとりができるようになります。

title: パイナップル
body: パイナップルは、南アメリカ原産の常緑性の熱帯果樹で、アブラナ科の植物です。果肉は黄色く、鮮やかな甘酸っぱい香りと味わいを持ち、豊富なビタミンCやカロテン、ポリフェノール、カルシウムなどの栄養素が含まれています。また、消化酵素であるブロメリンを含んでおり、食欲増進や消化促進効果があるとされています。切り方や調理法によっては、サラダやスムージー、パイやジュースなど、幅広い料理に使用されます。しかし、果肉とともに硬い芯があるため、適切な切り方をしなければ喉に詰まることがあるため注意が必要です。

title: 挑戦状
body: 挑戦状とは、自分や他人に対してある目標や困難を設定して、それに立ち向かうことを宣言する文書やメッセージのことです。ビジネスやスポーツ界でよく用いられ、自分自身や他人に向けたモチベーションやチャレンジ意識を高めるために発信されます。また、競技などで対戦相手に対して具体的な目標や条件を示して、対戦の勝敗を決定する場合にも用いられます。一般的には、挑戦状を提示した側が目標を達成すれば、設定された条件が満たされたことになります。しかし、目標を達成できなかった場合には、条件をクリアすることはできず、挑戦状の宣言者が敗北することになります。挑戦状を出すことで、自分自身のモチベーションアップや他者との競争意識の向上などが促されるため、自己成長や目標達成に向けた強い意志を持つことができます。

title: 成人
body: 成人とは、法律的には満20歳以上のことを指します。一般的には、心理的、社会的、経済的に独立し、自己責任で生活ができる人とも定義されます。成人になるためには、法律で定められた年齢に達するだけではなく、一定の法律的な資格が必要とされることがあります。たとえば、結婚や遺産相続、公的な契約の締結、選挙権や投票権などがあります。また、成人になるための手続きは国によって異なり、日本では20歳になった時点で自動的に成人となりますが、米国では成人になるためには18歳以上であることが必要です。成人となると、自己決定権や自己責任の重要性が増し、社会的義務や責任も負うことになります。

title: 焼き肉
body: 焼き肉は、肉をグリルやプレートで焼いて、熱い石焼などの上にのせたり、金属製のプレートに盛り付けて食べる日本の料理の一つです。焼く肉の種類には牛肉、豚肉、鶏肉などがあり、タレに漬けたり、特製のダレを付けたりして、味をつけます。また、野菜やキノコなども焼き肉と一緒に食べることができます。通常、家庭で楽しむことができるほか、専門の焼き肉店などでも提供されています。また、韓国や中国などでも似た料理があるため、アジア圏で広く親しまれています。焼き肉は、脂肪分やタンパク質、ビタミンB群といった栄養素を豊富に含み、食感や風味も楽しめるため、人気のある料理です。ただし、適度な量で楽しむことが重要で、高脂肪や高カロリーの肉を大量に食べると、健康に悪影響を与えることがあります。

title: 迷彩柄
body: 迷彩柄（めいさいがら）は、主に軍隊や警察などの職業用衣服に使用される柄です。主に緑色、茶色、灰色の組み合わせで構成され、自然環境に溶け込みやすいように設計されています。迷彩柄は、兵士や警察官が標的となることを防ぐために開発されたもので、敵が見つけづらく、かつ敵を発見しやすいという特徴があります。迷彩柄は現在ではファッションアイテムとしても一般的であり、スニーカーやジャケットなどのアイテムに使われています。ただし、軍事目的で使う場合は、規制があるため普通に販売される迷彩柄の服装品を国外に持ち出すのは厳禁であり、法律に抵触することもあります。

title: 竜巻
body: 竜巻とは、空中の気流が急激に回転し、地上に伸びる高速の渦巻状の気流現象です。竜巻は、雷雨の時に発生することが多く、空気の状態が不安定な場合に発生しやすいとされています。竜巻の強さは、F0からF5までの6段階に分類されます。F0は比較的弱い竜巻で、軽い被害しか出ませんが、F5は非常に強力な竜巻で、家屋や建物を巻き上げるなどの大規模な被害を引き起こすことがあります。竜巻が発生すると、突然の激しい風や大雨が降ってくるため、被害を受けないためには、速やかに建物の中に避難するか、安全な場所に逃げることが求められます。また、竜巻が発生した場合には、安全を確保するために、ニュースや天気予報などから最新の情報を収集するようにしましょう。